#include "occupancy_map.h"
#include <image_transport/image_transport.h>
#include <sensor_msgs/image_encodings.h>
#include <cv_bridge/CvBridge.h>
#include <opencv/cv.h>
#include <opencv/highgui.h>
#include <fstream>
#define DRAW_SIZE 1
#define DRAW_SCALE 1

namespace enc = sensor_msgs::image_encodings;

OccupancyMap::OccupancyMap():
  min_x_(0), min_y_(0), max_x_(0), max_y_(0), trajectory_list_(NULL) {
	CvScalar color;
	color.val[0] = 255;
	color.val[1] = 0;
	color.val[2] = 0;
	colors_.push_back(color);

	color.val[0] = 0;
	color.val[1] = 255;
	color.val[2] = 0;
	colors_.push_back(color);

	color.val[0] = 0;
	color.val[1] = 0;
	color.val[2] = 255;
	colors_.push_back(color);
}

int OccupancyMap::getGridXY(int x, int y) const {
	Coordinates pose(x, y);
	gridToEuclide(&pose);

	return getXY(pose.first, pose.second);
}

int OccupancyMap::getXY(int x, int y) const {
	Coordinates point(x, y);
	Map::const_iterator value = map_.find(point);

	if (value == map_.end())
		return UNKNOWN;

	return value->second;
}

void OccupancyMap::setXY(int x, int y, int value) {
	Coordinates point(x, y);
	map_[point] = value;

	if (x < min_x_)
		min_x_ = x;
	if (x > max_x_)
		max_x_ = x;
	if (y < min_y_)
		min_y_ = y;
	if (y > max_y_)
		max_y_ = y;
}

void OccupancyMap::clearXY(int x, int y) {
	Coordinates point(x, y);
	map_.erase(point);
}

int OccupancyMap::getSize(void) const {
	int val = abs(max_x_);

	if (abs(max_y_) > val)
		val = abs(max_y_);

	if (abs(min_x_) > val)
		val = abs(min_x_);

	if (abs(min_y_) > val)
		val = abs(min_y_);

	return (val * 2) + 1;
}

int OccupancyMap::getMinX(void) const {
	return min_x_;
}

int OccupancyMap::getMaxX(void) const {
	return max_x_;
}

int OccupancyMap::getMinY(void) const {
	return min_y_;
}

int OccupancyMap::getMaxY(void) const {
	return max_y_;
}

void OccupancyMap::addTrajectory(const Trajectory *trajectory) {
  trajectories_.push_back(trajectory);
}

void OccupancyMap::addTrajectoryList(const Trajectories *traj_list) {
  trajectory_list_ = traj_list;
}

void OccupancyMap::merge(const OccupancyMap& new_map) {
	const Map &new_map_h = new_map.getMapHandle();
  for (Map::const_iterator it = new_map_h.begin(); 
			 it != new_map_h.end(); ++it) {
		setXY(it->first.first, it->first.second, it->second);
	}
}

const Map& OccupancyMap::getMapHandle(void) const {
  return map_;
}

void OccupancyMap::euclideToGrid(Coordinates *pose) const {
	pose->first += getSize()/2;
	pose->second += getSize()/2;
}

void OccupancyMap::gridToEuclide(Coordinates *pose) const {
	pose->first -= getSize()/2;
	pose->second -= getSize()/2;
}

sensor_msgs::Image::Ptr OccupancyMap::getImgMessage(void) {
	IplImage* img = cvCreateImage(
			cvSize(getSize(), getSize()),
			IPL_DEPTH_8U, 3);

	CvScalar color_unknown;
	color_unknown.val[0] = GREY;
	color_unknown.val[1] = GREY;
	color_unknown.val[2] = GREY;

  for (int i = 0; i < getSize(); i++) {
    for (int j = 0; j < getSize(); j++) {	
			setPixel_(img, i, j, color_unknown);
		}
	}

	for (Map::iterator it = map_.begin(); it != map_.end(); ++it) {
    int value = it->second;
		Coordinates grid_point(it->first);
		euclideToGrid(&grid_point);

		int x = grid_point.first;
		int y = grid_point.second;
		CvScalar color;
		if (value == FREE) {
		  color.val[0] = WHITE;
		  color.val[1] = WHITE;
	  	color.val[2] = WHITE;
		} else {
		  color.val[0] = BLACK;
		  color.val[1] = BLACK;
	  	color.val[2] = BLACK;
		}
		setPixel_(img, x, y, color);
	}

	int i = 0;
	for (Trajectories::iterator traj_it = trajectories_.begin();
			traj_it != trajectories_.end(); ++traj_it, ++i) {
		CvScalar color = getColor_(i++);
		for(Trajectory::const_iterator pose_it = (*traj_it)->begin();
				pose_it != (*traj_it)->end(); ++pose_it){
			Coordinates grid_point(*pose_it);
			euclideToGrid(&grid_point);

			if (*pose_it == (*traj_it)->back()) {
				color.val[0] = 0;
				color.val[1] = 200;
				color.val[2] = 255;
			}
			drawPose_(grid_point, img, color);
		}
	}

	if (trajectory_list_) {
		for (Trajectories::const_iterator traj_it = trajectory_list_->begin();
				traj_it != trajectory_list_->end(); ++traj_it, ++i) {
			CvScalar color = getColor_(i++);
			for(Trajectory::const_iterator pose_it = (*traj_it)->begin();
					pose_it != (*traj_it)->end(); ++pose_it){
				Coordinates grid_point(*pose_it);
				euclideToGrid(&grid_point);

				if (*pose_it == (*traj_it)->back()) {
					color.val[0] = 0;
					color.val[1] = 200;
					color.val[2] = 255;
				}
				drawPose_(grid_point, img, color);
			}
		}
	}

	IplImage* img_scaled = cvCreateImage(
			cvSize(getSize() * DRAW_SCALE, getSize() * DRAW_SCALE),
			IPL_DEPTH_8U, 3);
	cvResize(img, img_scaled);

	sensor_msgs::Image::Ptr ptr = sensor_msgs::CvBridge::cvToImgMsg(img_scaled);
	cvReleaseImage(&img);
	cvReleaseImage(&img_scaled);
	return ptr;

}

void OccupancyMap::drawPose_(const Coordinates& pose, IplImage *img,
		const CvScalar& color) {
	int start_x = (pose.first - DRAW_SIZE);
	int start_y = (pose.second - DRAW_SIZE);
	int end_x = (pose.first + DRAW_SIZE);
	int end_y = (pose.second + DRAW_SIZE);
	if (start_x < 0)
		start_x = 0;
	if (start_y < 0)
		start_y = 0;
	if (end_x > img->width)
		end_x = img->width;
	if (end_y > img->height)
		end_y = img->height;
  ROS_DEBUG("sx %d sy %d ex %d ey %d", start_x, start_y, end_x, end_y);
	for(int i = start_x; i < end_x; i++){
		for(int j = start_y; j < end_y; j++){
			setPixel_(img, i, j, color);
		}
	}
}

void OccupancyMap::setPixel_(IplImage *img, int x, int y,
		const CvScalar &color) {
	((uchar *)(img->imageData + x*img->widthStep))[y*img->nChannels + 0] = 
		color.val[0];
	((uchar *)(img->imageData + x*img->widthStep))[y*img->nChannels + 1] = 
		color.val[1];
	((uchar *)(img->imageData + x*img->widthStep))[y*img->nChannels + 2] = 
		color.val[2];

}

const CvScalar OccupancyMap::getColor_(int idx) {
  return colors_[idx % colors_.size()];
}

